{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "307804a3-c02b-4a57-ac0d-172c30ddc851",
   "metadata": {},
   "source": [
    "# LlamaIndex + Pinecone "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99ba1864",
   "metadata": {},
   "source": [
    "In this tutorial, we show how to use LlamaIndex with Pinecone to answer complex queries over multiple data sources.  \n",
    "* While Pinecone provides a powerful and efficient retrieval engine,\n",
    "it remains challenging to answer complex questions that require multi-step reasoning and synthesis over many data sources.\n",
    "* With LlamaIndex, we combine the power of vector similiarty search and multi-step reasoning to delivery higher quality and richer responses.\n",
    "\n",
    "\n",
    "Here, we show 2 specific use-cases:\n",
    "1. compare and contrast queries over Wikipedia articles about different cities.\n",
    "2. temporal queries that require reasoning over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d48af8e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "#### Creating a Pinecone Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ce3143d-198c-4dd2-8e5a-c5cdf94f017a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/suo/miniconda3/envs/llama/lib/python3.9/site-packages/pinecone/index.py:4: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "import pinecone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "81fc463e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pinecone.init(environment=\"eu-west1-gcp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3500dd93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create index if it does not already exist\n",
    "# dimensions are for text-embedding-ada-002\n",
    "pinecone.create_index(\n",
    "    \"quickstart-index\", dimension=1536, metric=\"euclidean\", pod_type=\"p1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7b1544b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pinecone_index = pinecone.Index(\"quickstart-index\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "027b8cae",
   "metadata": {},
   "source": [
    "# Use-Case 1: Compare and Contrast"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26d56590",
   "metadata": {},
   "source": [
    "#### Load Dataset\n",
    "\n",
    "Fetch and load Wikipedia pages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c410eec0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "from llama_index import SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "583afcf9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "wiki_titles = [\n",
    "    \"Toronto\",\n",
    "    \"Seattle\",\n",
    "    \"San Francisco\",\n",
    "    \"Chicago\",\n",
    "    \"Boston\",\n",
    "    \"Washington, D.C.\",\n",
    "    \"Cambridge, Massachusetts\",\n",
    "    \"Houston\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "540a39b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import requests\n",
    "\n",
    "data_path = Path(\"data_wiki\")\n",
    "\n",
    "for title in wiki_titles:\n",
    "    response = requests.get(\n",
    "        \"https://en.wikipedia.org/w/api.php\",\n",
    "        params={\n",
    "            \"action\": \"query\",\n",
    "            \"format\": \"json\",\n",
    "            \"titles\": title,\n",
    "            \"prop\": \"extracts\",\n",
    "            \"explaintext\": True,\n",
    "        },\n",
    "    ).json()\n",
    "    page = next(iter(response[\"query\"][\"pages\"].values()))\n",
    "    wiki_text = page[\"extract\"]\n",
    "\n",
    "    if not data_path.exists():\n",
    "        Path.mkdir(data_path)\n",
    "\n",
    "    with open(data_path / f\"{title}.txt\", \"w\") as fp:\n",
    "        fp.write(wiki_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4b802084",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load all wiki documents\n",
    "city_docs = {}\n",
    "all_docs = []\n",
    "for wiki_title in wiki_titles:\n",
    "    city_docs[wiki_title] = SimpleDirectoryReader(\n",
    "        input_files=[data_path / f\"{wiki_title}.txt\"]\n",
    "    ).load_data()\n",
    "    all_docs.extend(city_docs[wiki_title])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee4473a-094f-4d0a-a825-e1213db07240",
   "metadata": {},
   "source": [
    "#### Build Indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0a2bcc07",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index import VectorStoreIndex, StorageContext\n",
    "from llama_index.vector_stores import PineconeVectorStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ba1558b3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building index for Toronto\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 20744 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 20744 tokens\n",
      "Building index for Seattle\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 16942 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 16942 tokens\n",
      "Building index for San Francisco\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 23837 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 23837 tokens\n",
      "Building index for Chicago\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 26082 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 26082 tokens\n",
      "Building index for Boston\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 18650 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 18650 tokens\n",
      "Building index for Washington, D.C.\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 21731 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 21731 tokens\n",
      "Building index for Cambridge, Massachusetts\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 12855 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 12855 tokens\n",
      "Building index for Houston\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 21845 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 21845 tokens\n"
     ]
    }
   ],
   "source": [
    "# Build index for each city document\n",
    "city_indices = {}\n",
    "index_summaries = {}\n",
    "for wiki_title in wiki_titles:\n",
    "    print(f\"Building index for {wiki_title}\")\n",
    "    # create storage context\n",
    "    vector_store = PineconeVectorStore(\n",
    "        pinecone_index=pinecone_index, namespace=wiki_title\n",
    "    )\n",
    "    storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "    # build index\n",
    "    city_indices[wiki_title] = VectorStoreIndex.from_documents(\n",
    "        city_docs[wiki_title], storage_context=storage_context\n",
    "    )\n",
    "\n",
    "    # set summary text for city\n",
    "    index_summaries[wiki_title] = f\"Wikipedia articles about {wiki_title}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04304299-fc3e-40a0-8600-f50c3292767e",
   "metadata": {},
   "source": [
    "#### Build Graph Query Engine for Compare & Contrast Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "35369eda",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.composability import ComposableGraph\n",
    "from llama_index.indices.keyword_table.simple_base import SimpleKeywordTableIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bedbb693-725f-478f-be26-fa7180ea38b2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "graph = ComposableGraph.from_indices(\n",
    "    SimpleKeywordTableIndex,\n",
    "    [index for _, index in city_indices.items()],\n",
    "    [summary for _, summary in index_summaries.items()],\n",
    "    max_keywords_per_chunk=50,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ba8ddcf0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.query.query_transform.base import DecomposeQueryTransform\n",
    "from llama_index.query_engine.transform_query_engine import TransformQueryEngine\n",
    "\n",
    "decompose_transform = DecomposeQueryTransform(verbose=True)\n",
    "\n",
    "custom_query_engines = {}\n",
    "for wiki_title in wiki_titles:\n",
    "    index = city_indices[wiki_title]\n",
    "    query_engine = index.as_query_engine()\n",
    "    query_engine = TransformQueryEngine(\n",
    "        query_engine,\n",
    "        query_transform=decompose_transform,\n",
    "        transform_extra_info={\"index_summary\": index_summaries[wiki_title]},\n",
    "    )\n",
    "    custom_query_engines[index.index_id] = query_engine\n",
    "\n",
    "custom_query_engines[graph.root_id] = graph.root_index.as_query_engine(\n",
    "    retriever_mode=\"simple\",\n",
    "    response_mode=\"tree_summarize\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "78235fbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# with query decomposition in subindices\n",
    "query_engine = graph.as_query_engine(custom_query_engines=custom_query_engines)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dd3616c",
   "metadata": {},
   "source": [
    "#### Run Compare & Contrast Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3f210a02",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.indices.keyword_table.retrievers:> Starting query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "> Starting query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['demographics', 'houston', 'contrast', 'seattle', 'toronto', 'compare']\n",
      "query keywords: ['demographics', 'houston', 'contrast', 'seattle', 'toronto', 'compare']\n",
      "INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: ['houston', 'seattle', 'toronto']\n",
      "> Extracted keywords: ['houston', 'seattle', 'toronto']\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Houston?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 7 tokens\n",
      "> [retrieve] Total embedding token usage: 7 tokens\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Houston?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1803 tokens\n",
      "> [get_response] Total LLM token usage: 1803 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Seattle?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 7 tokens\n",
      "> [retrieve] Total embedding token usage: 7 tokens\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Seattle?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1834 tokens\n",
      "> [get_response] Total LLM token usage: 1834 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Toronto?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 7 tokens\n",
      "> [retrieve] Total embedding token usage: 7 tokens\n",
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the demographics in Seattle, Houston, and Toronto.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query:  What is the population of Toronto?\n",
      "\u001b[0mINFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1909 tokens\n",
      "> [get_response] Total LLM token usage: 1909 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 233 tokens\n",
      "> [get_response] Total LLM token usage: 233 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 233 tokens\n",
      "> [get_response] Total LLM token usage: 233 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Compare and contrast the demographics in Seattle, Houston, and Toronto.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d24db9d4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: Seattle, Houston, and Toronto are all large cities\n",
      "with diverse populations. Houston is the largest of the three cities,\n",
      "with a population of 2,304,580 according to the 2020 U.S. census.\n",
      "Seattle is the second largest, with an estimated population of 730,000\n",
      "people. Toronto is the third largest, with a population of 6,202,225\n",
      "in 2021. All three cities have a diverse population, with a mix of\n",
      "different ethnicities, cultures, and religions. Houston is known for\n",
      "its large Hispanic population, while Seattle is known for its large\n",
      "Asian population. Toronto is known for its multiculturalism, with a\n",
      "large population of immigrants from all over the world.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.response.pprint_utils import pprint_response\n",
    "\n",
    "pprint_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a23412f",
   "metadata": {},
   "source": [
    "# Use-Case 2: Temporal Query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce276271",
   "metadata": {},
   "source": [
    "Temporal queries such as \"what happened after X\" is intuitive to humans, but can often confuse vector databases.  \n",
    "\n",
    "This is because the vector embedding will focus on the subject \"X\" rather than the imporant temporal cue. This results in irrelevant and misleading context that harms the final answer.  \n",
    "\n",
    "LlamaIndex solves this by explicitly maintainging node relationships and leverage LLM to automatically perform query expansion to find more relevant context.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3150f9a5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 20729 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 20729 tokens\n"
     ]
    }
   ],
   "source": [
    "from llama_index import SimpleDirectoryReader, StorageContext, VectorStoreIndex\n",
    "from llama_index.vector_stores import PineconeVectorStore\n",
    "\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../data/paul_graham\").load_data()\n",
    "\n",
    "# define storage context\n",
    "vector_store = PineconeVectorStore(\n",
    "    pinecone_index=pinecone_index, namespace=\"pg_essay_0.6.0\"\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "# build index\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    # override to store Node in document store in addition to vector store, necessary for the node postprocessor\n",
    "    store_nodes_override=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dbe01aa",
   "metadata": {},
   "source": [
    "We can define an auto prev/next node postprocessor to leverage LLM reasoning to help query expansion (with relevant additional nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "56f2d0d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.postprocessor.node import AutoPrevNextNodePostprocessor\n",
    "\n",
    "# define postprocessor\n",
    "node_postprocessor = AutoPrevNextNodePostprocessor(\n",
    "    docstore=index.storage_context.docstore,\n",
    "    service_context=index.service_context,\n",
    "    num_nodes=3,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "# define query engine\n",
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=1,\n",
    "    node_postprocessors=[node_postprocessor],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56896a41",
   "metadata": {},
   "source": [
    "#### Example 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ffd6c81a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 17 tokens\n",
      "> [retrieve] Total embedding token usage: 17 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1153 tokens\n",
      "> [get_response] Total LLM token usage: 1153 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1153 tokens\n",
      "> [get_response] Total LLM token usage: 1153 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "> Postprocessor Predicted mode: next\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 4204 tokens\n",
      "> [get_response] Total LLM token usage: 4204 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "# Infer that we need to search nodes after current one\n",
    "response = query_engine.query(\n",
    "    \"What did the author do after handing off Y Combinator to Sam Altman?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f91a50f9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: After handing off Y Combinator to Sam Altman, the\n",
      "author decided to take a break and pursue a completely different\n",
      "activity. He chose to paint and spent most of the rest of 2014\n",
      "painting. However, in November he ran out of steam and stopped\n",
      "painting. He then started writing essays again and wrote a few that\n",
      "weren't about startups. In March 2015, he started working on Lisp\n",
      "again, and spent the next four years writing a new Lisp called Bel in\n",
      "itself in Arc. He had to ban himself from writing essays during most\n",
      "of this time, or he would never have finished. In late 2015 he spent 3\n",
      "months writing essays, and when he went back to working on Bel he\n",
      "could barely understand the code.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.response.pprint_utils import pprint_response\n",
    "\n",
    "pprint_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22434b2e",
   "metadata": {},
   "source": [
    "In comparison, naive top-k retrieval results in irrelevant context and hallucinated answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a93d2432",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 17 tokens\n",
      "> [retrieve] Total embedding token usage: 17 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1028 tokens\n",
      "> [get_response] Total LLM token usage: 1028 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "# define query engine\n",
    "naive_query_engine = index.as_query_engine(\n",
    "    similarity_top_k=1,\n",
    ")\n",
    "\n",
    "response = naive_query_engine.query(\n",
    "    \"What did the author do after handing off Y Combinator to Sam Altman?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8ab35766-5ff2-4c4f-98e6-81039d99cb44",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: After handing off Y Combinator to Sam Altman, the\n",
      "author went on to found OpenAI, a research laboratory dedicated to\n",
      "artificial intelligence. He also wrote a book, \"The Launch Pad: Inside\n",
      "Y Combinator, Silicon Valley's Most Exclusive School for Startups,\"\n",
      "and became a partner at Founders Fund, a venture capital firm.\n",
      "______________________________________________________________________\n",
      "Source Node 1/1\n",
      "Document ID: 204a1bd3-95cd-4421-accb-469fbd876a00\n",
      "Similarity: 0.839429557\n",
      "Text: in. We also noticed that the startups were becoming one\n",
      "another's customers. We used to refer jokingly to the \"YC GDP,\" but as\n",
      "YC grows this becomes less and less of a joke. Now lots of startups\n",
      "get their initial set of customers almost entirely from among their\n",
      "batchmates.  I had not originally intended YC to be a full-time job. I\n",
      "was going to ...\n"
     ]
    }
   ],
   "source": [
    "pprint_response(response, show_source=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64117464",
   "metadata": {},
   "source": [
    "#### Example 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "221001e4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 17 tokens\n",
      "> [retrieve] Total embedding token usage: 17 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1131 tokens\n",
      "> [get_response] Total LLM token usage: 1131 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 1131 tokens\n",
      "> [get_response] Total LLM token usage: 1131 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "> Postprocessor Predicted mode: previous\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 4414 tokens\n",
      "> [get_response] Total LLM token usage: 4414 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "# Infer that we need to search nodes before current one\n",
    "response = query_engine.query(\n",
    "    \"What did the author do before handing off Y Combinator to Sam Altman?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "a6bf950f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: Before handing off Y Combinator to Sam Altman, the\n",
      "author worked on several different projects. He wrote essays and\n",
      "published them online, worked on spam filters, painted, cooked for\n",
      "groups, bought a building in Cambridge, and started Y Combinator. He\n",
      "also gave a talk at a Lisp conference, wrote a postscript file of the\n",
      "talk and posted it online, and started angel investing. He also\n",
      "schemed with Robert and Trevor about projects they could work on\n",
      "together, and then he and Jessica Livingston started their own\n",
      "investment firm. They created the Summer Founders Program, which was a\n",
      "summer program for undergrads to start startups instead of getting\n",
      "temporary jobs at tech companies. They also created the batch model,\n",
      "which was to fund a bunch of startups all at once, twice a year, and\n",
      "then to spend three months focusing intensively on trying to help\n",
      "them. He also worked on a new version of Arc with Robert, wrote to\n",
      "test the new Arc, and noticed the advantages of scale as YC grew.\n",
      "______________________________________________________________________\n",
      "Source Node 1/4\n",
      "Document ID: 2f59abe3-aa5c-4204-aaac-128097db0022\n",
      "Similarity: None\n",
      "Text: of the most conspicuous patterns I've noticed in my life is how\n",
      "well it has worked, for me at least, to work on things that weren't\n",
      "prestigious. Still life has always been the least prestigious form of\n",
      "painting. Viaweb and Y Combinator both seemed lame when we started\n",
      "them. I still get the glassy eye from strangers when they ask what I'm\n",
      "writing...\n",
      "______________________________________________________________________\n",
      "Source Node 2/4\n",
      "Document ID: 50f7af32-d191-4793-b4d6-beb97f59886c\n",
      "Similarity: None\n",
      "Text: I doing this? If this vision had to be realized as a company,\n",
      "then screw the vision. I'd build a subset that could be done as an\n",
      "open source project.  Much to my surprise, the time I spent working on\n",
      "this stuff was not wasted after all. After we started Y Combinator, I\n",
      "would often encounter startups working on parts of this new\n",
      "architecture, and...\n",
      "______________________________________________________________________\n",
      "Source Node 3/4\n",
      "Document ID: 86508708-65ec-4b1a-a666-ed858d60e953\n",
      "Similarity: None\n",
      "Text: start a startup. Maybe they'd be able to avoid the worst of the\n",
      "mistakes we'd made.  So I gave this talk, in the course of which I\n",
      "told them that the best sources of seed funding were successful\n",
      "startup founders, because then they'd be sources of advice too.\n",
      "Whereupon it seemed they were all looking expectantly at me. Horrified\n",
      "at the prospect o...\n",
      "______________________________________________________________________\n",
      "Source Node 4/4\n",
      "Document ID: f1bf5f95-afc8-403b-ae20-af4c1e011b99\n",
      "Similarity: 0.848301291\n",
      "Text: due to our ignorance about investing. We needed to get\n",
      "experience as investors. What better way, we thought, than to fund a\n",
      "whole bunch of startups at once? We knew undergrads got temporary jobs\n",
      "at tech companies during the summer. Why not organize a summer program\n",
      "where they'd start startups instead? We wouldn't feel guilty for being\n",
      "in a sense...\n"
     ]
    }
   ],
   "source": [
    "pprint_response(response, show_source=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
